import base64
import json
import uuid

import requests
import base64
import logging
import time
from openai import OpenAI

from google.oauth2 import service_account
from google.auth.transport.requests import Request
import sys, os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from utils import sr_utils, sr_s3
from utils.sr_utils import delete_temp_files, get_video_info

import agent_config


def save_base64_video(b64_data, output_file="output_video.mp4"):
    try:
        # 解码 Base64 数据
        video_data = base64.b64decode(b64_data)
        # 写入文件
        with open(output_file, "wb") as f:
            f.write(video_data)

        print(f"视频文件已成功保存为: {os.path.abspath(output_file)}")
        return os.path.abspath(output_file)
    except Exception as e:
        print(f"保存视频失败: {e}")
        return None


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 配置
DEFAULT_PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "aigc-veo")
GOOGLE_APPLICATION_CREDENTIALS = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
print("GOOGLE_APPLICATION_CREDENTIALS:", GOOGLE_APPLICATION_CREDENTIALS)

DEFAULT_LOCATION = "us-central1"
MODEL_IDS = {
    "veo-2.0": "veo-2.0-generate-001",
    "veo-3.0": "veo-3.0-generate-preview",
    "veo-3.1": "veo-3.1-generate-preview"
}
models_with_audio_support = {
    "veo-3.0-generate-001",
    "veo-3.0-generate-preview",
    "veo-3.1-generate-preview"
}

# 模型能力配置
MODEL_CAPABILITIES = {
    "veo-2.0-generate-001": {
        "support_text_to_video": True,
        "support_image_to_video": True,
        "support_keyframe": True,
        "durations": [5, 6, 7, 8],
        "aspect_ratios": ["16:9", "9:16"],
        "image_to_video_durations": [5, 6, 7, 8],
        "image_to_video_aspect_ratios": ["16:9", "9:16"]
    },
    "veo-3.0-generate-preview": {
        "support_text_to_video": True,
        "support_image_to_video": True,
        "support_keyframe": True,
        "durations": [4, 6, 8],  # ⚠️ 只支持 4/6/8 秒
        "aspect_ratios": ["16:9", "9:16"],
        "image_to_video_durations": [4, 6, 8],
        "image_to_video_aspect_ratios": ["16:9", "9:16"]
    },
    "veo-3.1-generate-preview": {
        "support_text_to_video": True,
        "support_image_to_video": True,
        "support_keyframe": True,
        "durations": [4, 6, 8],  # ⚠️ 只支持 4/6/8 秒
        "aspect_ratios": ["16:9", "9:16"],
        "image_to_video_durations": [8],  # ⚠️ 图生视频只支持8秒
        "image_to_video_aspect_ratios": ["16:9"]  # ⚠️ 图生视频只支持16:9
    }
}


def validate_model_params(model_id, model_type, duration, aspect_ratio):
    """
    验证模型参数是否符合要求
    
    Args:
        model_id: 模型ID (如 veo-3.1-generate-preview)
        model_type: 模型类型 (1=文生, 2/3=图生首帧, 4=首尾帧)
        duration: 视频时长（秒）
        aspect_ratio: 宽高比 (16:9, 9:16)
    
    Returns:
        (bool, str): (是否通过, 错误信息)
    """
    if model_id not in MODEL_CAPABILITIES:
        return False, f"不支持的模型: {model_id}"
    
    caps = MODEL_CAPABILITIES[model_id]
    
    # 验证 modelType
    if model_type == 1 and not caps["support_text_to_video"]:
        return False, f"{model_id} 不支持文生视频"
    if model_type in [2, 3] and not caps["support_image_to_video"]:
        return False, f"{model_id} 不支持图生视频"
    if model_type == 4 and not caps["support_keyframe"]:
        return False, f"{model_id} 不支持首尾帧模式"
    
    # 验证时长
    if duration not in caps["durations"]:
        return False, f"{model_id} 不支持 {duration}秒，支持的时长: {caps['durations']}"
    
    # 图生视频特殊验证 (modelType 2, 3, 4)
    if model_type in [2, 3, 4]:
        if duration not in caps["image_to_video_durations"]:
            return False, f"{model_id} 图生视频不支持 {duration}秒，支持的时长: {caps['image_to_video_durations']}"
        
        if aspect_ratio not in caps["image_to_video_aspect_ratios"]:
            return False, f"{model_id} 图生视频不支持 {aspect_ratio}，支持的宽高比: {caps['image_to_video_aspect_ratios']}"
    else:
        # 文生视频验证宽高比
        if aspect_ratio not in caps["aspect_ratios"]:
            return False, f"{model_id} 不支持 {aspect_ratio}，支持的宽高比: {caps['aspect_ratios']}"
    
    return True, ""


def download_single_image(img_url, storage_service, bucket_name, save_dir="./upload"):
    """下载单张图片到本地并返回路径"""
    try:
        os.makedirs(save_dir, exist_ok=True)
        name = str(uuid.uuid4())
        ext = img_url[img_url.rfind('.'):] if '.' in img_url else '.jpg'
        save_path = os.path.join(save_dir, f"{name}{ext}")

        # 调用S3工具下载图片
        success, err = sr_s3.download_file(
            img_url, save_path, serviceName=storage_service, bucketName=bucket_name
        )
        if not success:
            return None, f"下载失败: {err}"
        # 压缩图片（确保符合模型要求）
        # success, new_path = sr_utils.video_compress_image_size_path(save_path, 25, (200, 6000), (200, 6000))
        # if not success:
        #     delete_temp_files(new_path)
        #     return None, "图片压缩失败" + new_path
        # return new_path, None
        return save_path, None
    except Exception as e:
        return None, f"处理图片异常: {str(e)}"


# 函数：编码图片为 base64
def encode_image(image_path):
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except FileNotFoundError:
        logger.error(f"图片文件未找到: {image_path}")
        raise
    except Exception as e:
        logger.error(f"编码图片失败: {e}")
        raise


# 函数：获取访问令牌
def get_access_token():
    try:
        credentials = service_account.Credentials.from_service_account_file(
            GOOGLE_APPLICATION_CREDENTIALS,
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )

        # 如果 token 过期，刷新它
        if not credentials.token or credentials.expired:
            credentials.refresh(Request())

        print("获取访问令牌成功:")
        return credentials.token
    except Exception as e:
        print(f"获取访问令牌失败: {e}")
        raise


# 函数：轮询操作状态
def poll_operation(operation_name, model_id, project_id=DEFAULT_PROJECT_ID, location=DEFAULT_LOCATION, timeout=1800, poll_interval=15):
    """
    轮询操作状态直到完成或超时（使用新的 POST 接口）。

    Args:
        operation_name (str): 操作的名称（从 generate_video 返回）。
        project_id (str): Google Cloud 项目 ID。
        location (str): API 位置。
        timeout (int): 最大等待时间（秒）。
        poll_interval (int): 轮询间隔时间（秒）。

    Returns:
        dict: 操作的最终结果。
    """
    # 构造 API 端点
    api_endpoint = f"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:fetchPredictOperation"

    # 获取访问令牌
    access_token = get_access_token()
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json"
    }

    # 请求体
    payload = {
        "operationName": operation_name
    }

    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            response = requests.post(api_endpoint, headers=headers, json=payload)
            response.raise_for_status()
            result = response.json()
            if result.get("done"):
                # 保存结果
                output_dir = "./poll_results"
                os.makedirs(output_dir, exist_ok=True)
                timestamp = int(time.time())
                output_path = os.path.join(output_dir, f"poll_result_{timestamp}.txt")
                with open(output_path, "w", encoding="utf-8") as f:
                    f.write(json.dumps(result, indent=2, ensure_ascii=False))
                print(f"轮询结果已写入")
                # 保存结果

                print("操作已完成")
                logger.info("操作已完成")
                if "error" in result:
                    # 在这里处理错误，而不是在调用函数中处理
                    error_result = handle_video_generation_error(result)
                    if error_result:
                        raise Exception(error_result)
                    raise Exception(f"操作失败: {result['error']['message']}")
                return result.get("response", {})
            else:
                logger.info(f"操作 {operation_name} 仍在进行中...")
                time.sleep(poll_interval)
        except requests.exceptions.HTTPError as http_err:
            logger.error(f"轮询 HTTP 错误: {http_err}")
            logger.error(f"响应内容: {response.text}")
            raise
        except Exception as e:
            logger.error(f"轮询失败: {e}")
            raise

    logger.error(f"操作 {operation_name} 超时")
    raise TimeoutError(f"操作 {operation_name} 在 {timeout} 秒内未完成")


# 函数：保存视频文件
def download_from_gcs(gcs_uri, output_file):
    """
    从 Google Cloud Storage 下载视频文件。

    Args:
        gcs_uri (str): GS:// 格式的视频 URI。
        output_file (str): 输出视频文件路径。
    """
    from google.cloud import storage

    try:
        # 解析 GCS URI (格式: gs://bucket/object)
        parts = gcs_uri.replace("gs://", "").split("/")
        bucket_name = parts[0]
        object_name = "/".join(parts[1:])

        # 初始化存储客户端
        storage_client = storage.Client()
        bucket = storage_client.bucket(bucket_name)
        blob = bucket.blob(object_name)

        # 下载文件
        blob.download_to_filename(output_file)
        logger.info(f"视频已从 GCS 下载到: {output_file}")
    except Exception as e:
        logger.error(f"从 GCS 下载视频失败: {e}")
        raise


def save_video(video_data, output_file):
    """
    兼容旧版: 将 Base64 编码的视频数据保存到文件。

    Args:
        video_data (str): Base64 编码的视频数据。
        output_file (str): 输出视频文件路径。
    """
    try:
        with open(output_file, "wb") as f:
            f.write(base64.b64decode(video_data))
        logger.info(f"视频已保存到: {output_file}")
    except Exception as e:
        logger.error(f"保存视频失败: {e}")
        raise


# 新增函数：处理视频生成错误
def handle_video_generation_error(result):
    """
    处理视频生成错误，识别特定的违规代码并返回格式化的错误信息

    Args:
        result (dict): 包含错误信息的操作结果

    Returns:
        dict: 格式化的错误信息，如果不是特定错误则返回 None
    """
    if "error" not in result:
        return None

    error_code = result["error"].get("code", 0)
    error_message = result["error"].get("message", "")

    # 如果是违反负责任 AI 实践的错误
    if error_code == 3:
        print(f"视频生成失败: {error_message}")

        # 定义需要检测的违规代码
        violation_codes = ["90789179", "63429089", "43188360", "78610348", "61493863", "56562880", "32635315",
                           "58061214"]

        # 检查错误消息中是否包含特定的违规代码
        found_code = None
        for code in violation_codes:
            if code in error_message:
                found_code = code
                print(f"找到违规代码: {found_code}")
                break

        # 根据代码确定违规类型
        violation_type_map = {
            "43188360": "色情内容",
            "63429089": "色情内容",
            "90789179": "色情内容",
            "78610348": "有害内容",
            "61493863": "暴力",
            "56562880": "暴力",
            "32635315": "粗俗",
            "58061214": "拒绝生成描绘儿童的内容"
        }

        if found_code and found_code in violation_type_map:
            violation_type = violation_type_map[found_code]
            error_msg = f"视频生成失败: 检测到{violation_type}({found_code})，违反Google的负责任AI实践"
            status = 7
            code = 7202
        else:
            error_msg = f"视频生成失败: {error_message}"
            status = 2
            code = 6001

        return {
            "status": status,
            "width": 0,
            "height": 0,
            "url": None,
            "msg": error_msg,
            "code": code
        }

    return None


# 函数：生成视频
def generate_video(model="veo-2.0", generateAudio=False, input_type="text", prompt="", image_path=None, output_file="output_video.mp4", project_id=DEFAULT_PROJECT_ID, location=DEFAULT_LOCATION, durationSeconds=5, aspectRatio="16:9", timeout=1800, poll_interval=20):
    # 验证参数
    if input_type not in ["text", "image"]:
        logger.error(f"无效的输入类型: {input_type}，必须为 'text' 或 'image'")
        raise ValueError("输入类型必须为 'text' 或 'image'")
    if input_type == "image" and not image_path:
        logger.error("当输入类型为图片时，必须提供 image_path")
        raise ValueError("image_path 不能为空")

    if not prompt:
        logger.error("提示文本不能为空")
        raise ValueError("提示文本不能为空")

    # 获取模型 ID
    model_id = model

    # 构造 API 端点
    print("project_id:", project_id)
    api_endpoint = f"https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers/google/models/{model_id}:predictLongRunning"

    # 准备请求负载
    instance = {"prompt": prompt}
    if input_type == "image" and image_path:
        if isinstance(image_path, str):
            # 单图路径："cat.jpg"
            encoded_image = encode_image(image_path)
            instance["image"] = {
                "bytesBase64Encoded": encoded_image,
                "mimeType": "image/jpeg"
            }
        elif isinstance(image_path, list):
            if len(image_path) == 0:
                raise ValueError("image_path 列表为空，至少需要一个图片路径")

            encoded_image = encode_image(image_path[0])
            instance["image"] = {
                "bytesBase64Encoded": encoded_image,
                "mimeType": "image/jpeg"
            }

            if len(image_path) >= 2:
                # 多图情况：提供 lastFrame
                encoded_last_image = encode_image(image_path[-1])
                instance["lastFrame"] = {
                    "bytesBase64Encoded": encoded_last_image,
                    "mimeType": "image/jpeg"
                }
            # 如果是 ["cat.jpg"] 这种单图列表，则不加 lastFrame

    payload = {
        "instances": [instance],
        "parameters": {
            "personGeneration": "allow_all",
            "aspectRatio": aspectRatio,  # 可选: "16:9" 或 "9:16"
            "durationSeconds": int(durationSeconds),  # 视频长度 (5-8 秒)
            "sampleCount": 1,  # 生成视频数量
        }
    }
    if generateAudio:
        payload["parameters"]["generateAudio"] = True

    print(payload)
    # 获取访问令牌
    access_token = get_access_token()

    # 设置请求头
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json"
    }

    # 在 generate_video 函数中的适当位置（替换原有代码）
    try:
        # 发送 POST 请求到 Vertex AI Veo API
        logger.info(f"向 {api_endpoint} 发送请求...")
        response = requests.post(api_endpoint, headers=headers, json=payload)
        response.raise_for_status()

        # 解析响应
        operation = response.json()
        operation_name = operation.get("name")
        logger.info(f"操作已开始: {operation_name}")

        # 轮询操作状态
        result = poll_operation(operation_name, model_id, project_id, location, timeout=timeout, poll_interval=poll_interval)

        # 提取 Base64 视频数据
        b64_video = result["videos"][0]["bytesBase64Encoded"]

        # 保存为视频文件
        output_file_path = save_base64_video(b64_video, output_file)

        return output_file_path

    except requests.exceptions.HTTPError as http_err:
        logger.error(f"HTTP 错误: {http_err}")
        logger.error(f"响应内容: {response.text}")
        raise
    except Exception as e:
        logger.error(f"生成视频失败: {e}")
        # 检查异常是否包含我们定义的错误格式
        if hasattr(e, 'args') and len(e.args) > 0:
            error_arg = e.args[0]
            # 如果异常参数是一个字典，并且包含 status 字段，说明是已经处理过的错误
            if isinstance(error_arg, dict) and 'status' in error_arg:
                return error_arg

        # 否则返回通用错误
        error_msg = str(e)
        return {
            "status": 2,
            "width": 0,
            "height": 0,
            "url": None,
            "msg": f"视频生成失败: {error_msg}",
            "code": 6001
        }


# 示例调用
def ViduGenerator_google_veo(body):
    """统一处理视频生成任务（入口函数）"""

    task_id = body.get('taskId', 'unknown')
    # queue_result = body.get('QueueResult')
    storage_service = body.get('storageService', 's3')
    bucket_name = body.get('bucketName')
    s3_storage = body.get('storedPrefix')

    # 必传参数校验
    prompt = body.get('prompt')
    duration = body.get('duration', 5)
    ratio = body.get('ratio', "16:9")
    bgm = body.get('bgm', False)

    # 获取模型的轮询结果配置
    config_poll_time = agent_config.get_poll_config("google")
    timeout = config_poll_time.get('timeout', 600)
    poll_interval = config_poll_time.get('poll_interval', 20)

    first_image_url = body.get('first_image', None)
    last_image_url = body.get('last_image', None)

    storageService = body.get('storageService', 's3')
    bucketName = body.get('bucketName')

    # 模型类型
    model_type = body.get("modelType", 0)
    # 输出路径
    name = str(uuid.uuid4())
    output_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'results', name))
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f"{task_id}_video.mp4")
    moduleName = body.get('moduleName')
    
    # ✅ 参数验证：检查模型是否支持当前配置
    is_valid, error_msg = validate_model_params(moduleName, model_type, int(duration), ratio)
    if not is_valid:
        logger.error(f"参数验证失败: {error_msg}")
        return {
            "status": 0,
            "width": 0,
            "height": 0,
            "url": None,
            "msg": error_msg,
            "code": 6003  # 参数错误
        }
    # 下载图片
    image_urls = []
    image_path = first_image_url
    input_image = None
    # 文字转视频示例
    try:
        # 文生视频
        if model_type == 1:
            input_type = "text"
            input_image = None
        elif model_type == 2:
            # 图生视频
            input_type = "image"
            if not image_path:
                return {"status": 2, "width": 0, "height": 0, "url": None, "msg": "必须提供图片", 'code': 6003}

            if first_image_url and first_image_url.strip():
                image_path, err = download_single_image(first_image_url, storage_service, bucket_name)
                image_urls.append(image_path)
                input_image = image_urls

            # 图片是否存在；
            if not image_path:
                return {"status": 2, "width": 0, "height": 0, "url": None, "msg": "图片下载失败", 'code': 6005}
        elif model_type == 3:
            input_type = "image"
            # 图生视频依据第一帧图片
            if not image_path:
                return {"status": 2, "width": 0, "height": 0, "url": None, "msg": "必须提供图片", 'code': 6003}

            if first_image_url and first_image_url.strip():
                image_path, err = download_single_image(first_image_url, storage_service, bucket_name)
                image_urls.append(image_path)
                input_image = image_urls

            # 图片是否存在；
            if not image_path:
                return {"status": 2, "width": 0, "height": 0, "url": None, "msg": "图片下载失败", 'code': 6003}
        elif model_type == 4:
            # 图生视频依据尾帧图片
            if not first_image_url or not last_image_url:
                return {"status": 2, "width": 0, "height": 0, "url": None, "msg": "必须提供首尾图片", 'code': 6005}

            if first_image_url and first_image_url.strip():
                image_path, err = download_single_image(first_image_url, storage_service, bucket_name)
                image_urls.append(image_path)
            if last_image_url and last_image_url.strip():
                image_path, err = download_single_image(last_image_url, storage_service, bucket_name)
                image_urls.append(image_path)

            input_type = "image"
            input_image = image_urls
            # 图片是否存在；
            if not image_path:
                return {"status": 2, "width": 0, "height": 0, "url": None, "msg": "图片下载失败", 'code': 6003}

        # 根据模型是否支持音频以及bgm参数是否为True来决定是否启用音频生成
        generateAudio = (moduleName in models_with_audio_support) and bgm

        local_file = generate_video(
            model=moduleName,
            generateAudio=generateAudio,
            input_type=input_type,
            prompt=prompt,
            image_path=input_image,
            output_file=output_file,
            project_id=DEFAULT_PROJECT_ID,
            location=DEFAULT_LOCATION,
            durationSeconds=duration,
            aspectRatio=ratio,
            timeout=timeout,
            poll_interval=poll_interval
        )

        # 检查返回结果是否为字典（即错误状态）
        if isinstance(local_file, dict) and local_file.get("status") == 7:
            return local_file  # 直接返回错误状态

        # 3. 上传视频
        print(f"开始上传视频: {local_file}")
        print(s3_storage)
        s3_video_key = sr_s3.upload_login_or_visitor(local_file, s3_storage, serviceName=storage_service, bucketName=bucket_name)
        if not s3_video_key:
            return False, f"视频上传r2失败: {err}"

        logger.info(f"文字转视频任务已完成，操作: {local_file}")
        # 4. 获取视频信息
        local_video = get_video_info(local_file)
        return {
            "status": 1,
            "width": local_video["width"],
            "height": local_video["height"],
            'local_path': local_file,
            "model_mode": "std",  # 模型类型标准模式
            "url": s3_video_key,
            'msg': "success",
            "code": ""
        }
    except Exception as e:
        logger.error(f"视频失败: {e}")
        return {
            "status": 2,
            "width": 0,
            "height": 0,
            "url": None
        }


if __name__ == "__main__":
    # 示例任务参数v
    body = {
        "taskId": "test_001",
        "moduleName": "veo-3.0-generate-preview",
        "prompt": "一个兔子在跳舞",
        "modelType": 1,  # 文字生视频
        # "first_image": "video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1752733793325.jpg",
        # "last_image": "video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1752733793325.jpg",
        "duration": "5",
        "storageService": "r2",
        "bucketName": "picwand",
        "storedPrefix": "video/visitor/image_to_video",
        "QueueResult": "AILocalVideoResult"
    }
    # body = {'taskId': 'a01ab3fcc98346c9d630', 'prompt': 'A ship was sailing on the sea when it suddenly hit an iceberg.', 'first_image': '', 'last_image': '', 'duration': '5', 'fps': '24', 'resolution': '720p', 'ratio': '16:9', 'appId': '1', 'taskType': 'pwai_video_itf',
    #         'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 4, 'moduleType': 23, 'moduleName': 'veo-2.0-generate-001', 'mResolution': '720p',
    #         'mode': 'std', 'negative_prompt': '', 'cfg_scale': 0}
    # 初始化视频生成器
    result = ViduGenerator_google_veo(body)
    print(result)
